import xarray as xr
import numpy as np

def simple_correct_average():
    """
    简单修正：按位置对应求平均
    """
    tynames = ['mojie', 'dusurui', 'gaemi', 'haikui', 'kangni', 'shantuo', 'saola', 'koinu']
    
    # 读取所有数据
    all_data = []
    for tyname in tynames:
        file_path = f'/data/gsj/typlot/typlot/aligned/aligned_{tyname}_pmin_umax_dist_extended.nc'
        ds = xr.open_dataset(file_path)
        all_data.append(ds)
    
    # 使用第一个文件的结构创建模板
    template = all_data[0]
    
    # 手动计算平均（按位置对应）
    result_ds = template.copy()
    
    for var_name in template.data_vars:
        if var_name not in template.dims:
            # 堆叠所有数据
            stacked = np.stack([ds[var_name].values for ds in all_data], axis=0)
            # 计算平均，忽略NaN
            result_ds[var_name].values = np.nanmean(stacked, axis=0)
    
    # 使用统一的start_time（相对时间索引）
    result_ds = result_ds.assign_coords({
        'start_time': [str(i) for i in range(57)]
    })
    
    # 保存
    output_file = '/data/gsj/typlot/typlot/aligned/average_8typhoons.nc'
    result_ds.to_netcdf(output_file)
    print(f"简单修正的平均文件已保存: {output_file}")
    print(f"维度: {dict(result_ds.sizes)}")
    
    return result_ds

# 执行简单修正
simple_avg_ds = simple_correct_average()